import torch
from torch import nn

# L = 5 , N = 32 , D = 10 ,H = 200 , NL = 4
# D,H,NL
rnn = nn.GRU(10, 200, 4)
# L,N,D
# NL,N,H
input = torch.randn(5, 32, 10)
h0 = torch.randn(4, 32, 200)
output, hn = rnn(input, h0)
# 5, 32, 200
# L,N,H
print(output.shape)
# NL,N,H
print(hn.shape)

# L = 20 , N = 32 , D = 10 ,H = 256 , NL = 4
# D,H,NL
rnn = nn.LSTM(10, 256, 4)
# L,N,D
input = torch.randn(20, 32, 10)
# NL,N,H
h0 = torch.randn(4, 32, 256)
# NL,N,H
c0 = torch.randn(4, 32, 256)
# NL,N,H
output, (hn, cn) = rnn(input, (h0, c0))
# L,N,H
print(output.shape)
# NL,N,H
print(hn.shape)
# NL,N,H
print(cn.shape)

# L = 5 , N = 32 , D = 10 ,H = 200 , NL = 4
# D,H,NL
rnn = nn.GRU(10, 200, 4, bidirectional=True)
# L,N,D
# NL,N,H
input = torch.randn(5, 32, 10)
h0 = torch.randn(8, 32, 200)
output, hn = rnn(input, h0)
# 5, 32, 200
# L,N,H*2
print(output.shape)
# NL,N,H
print(hn.shape)
